import mlflow
import pandas as pd

def generate_recommendations_table_with_best(experiment_ids, group_type="sim"):
    all_rows = []

    for exp_id in experiment_ids:
        runs = mlflow.search_runs(
            experiment_ids=[exp_id],
            output_format="list"
        )
        for run in runs:
            if run.data.params.get("group_set") != 'test' or run.data.params.get("group_type") != group_type:
                continue

            dataset = run.data.params.get("dataset", f"Exp-{exp_id}")
            approach = run.data.params.get("recommender_strategy", None)
            aggregation = run.data.params.get("SAE_fusion_strategy", None)

            row_key = (approach, aggregation)
            metrics = {
                (dataset, "G/mean"): run.data.metrics.get("CommonItemsNDCG20/median"),
                (dataset, "U/mean"): run.data.metrics.get("NDCG20/mean"),
                (dataset, "U/min"): run.data.metrics.get("NDCG20/min"),
                (dataset, "Pop"): run.data.metrics.get("Popularity/mean"),
            }

            all_rows.append((row_key, metrics))

    # Build DataFrame from records
    records = {}
    for key, metrics in all_rows:
        if key not in records:
            records[key] = {}
        records[key].update(metrics)

    df = pd.DataFrame.from_dict(records, orient="index")
    df.index.names = ["Approach", "Aggregation"]

    # Sort and reindex columns by dataset then metric
    df = df.sort_index(axis=1, level=[0, 1]).sort_values(
        by=("MovieLens", "G/mean"), ascending=False
    ).reset_index()

    return df.reset_index().round(3)

experiment_ids = ['523100174176986081', '333391697323445885']

df = generate_recommendations_table_with_best(experiment_ids, group_type="sim")
df

experiment_ids = ['523100174176986081', '333391697323445885']

df = generate_recommendations_table_with_best(experiment_ids, group_type="random")
df

experiment_ids = ['523100174176986081', '333391697323445885']

df = generate_recommendations_table_with_best(experiment_ids, group_type="div")
df